import torch
import random
import numpy as np
import torch.nn as nn
import torch.nn.functional as F
import torch.optim as optim
from collections import deque

LR_ACTOR = 1e-4
LR_CRITIC = 1e-4
GAMMA = 0.95
MEMORY_SIZE = 4096
BATCH_SIZE = 128
TAU = 1e-3
HIDDEN_DIM = 128
GRAD_CLIP = 0.5
ENTROPY_DECAY = 0.997

device = torch.device("cuda" if torch.cuda.is_available() else "cpu")


class Actor(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim=HIDDEN_DIM):
        super(Actor, self).__init__()
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, hidden_dim)
        self.fc4 = nn.Linear(hidden_dim, hidden_dim)
        self.fc5 = nn.Linear(hidden_dim, action_dim)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        x = torch.relu(self.fc3(x))
        x = torch.relu(self.fc4(x))
        x = torch.softmax(self.fc5(x), dim=-1)
        return x


class QNetwork(nn.Module):
    def __init__(self, state_dim, action_dim, hidden_dim=HIDDEN_DIM):
        super(QNetwork, self).__init__()
        self.fc1 = nn.Linear(state_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc2 = nn.Linear(hidden_dim, hidden_dim)
        self.fc3 = nn.Linear(hidden_dim, action_dim)

    def forward(self, x):
        x = torch.relu(self.fc1(x))
        x = torch.relu(self.fc2(x))
        return self.fc3(x)


class ReplayMemory:
    def __init__(self, capacity):
        self.buffer = deque(maxlen=capacity)
        self.priorities = deque(maxlen=capacity)
        self.prio_alpha = 0.6
        self.prio_eps = 1e-6

    def add_memo(self, state, action, reward, next_state):
        state = np.expand_dims(state, 0)
        next_state = np.expand_dims(next_state, 0)
        self.buffer.append((state, action, reward, next_state))
        if len(self.priorities) == 0:
            self.priorities.append(1.0)
        else:
            self.priorities.append(max(self.priorities))

    def sample(self, batch_size):
        n = len(self.buffer)
        prios = np.array(self.priorities, dtype=np.float32)
        probs = (prios + self.prio_eps) ** self.prio_alpha
        probs = probs / probs.sum()
        idxs = np.random.choice(n, size=batch_size, replace=False, p=probs)
        buf_list = list(self.buffer)
        samples = [buf_list[i] for i in idxs]
        state, action, reward, next_state = zip(*samples)
        return np.concatenate(state), action, reward, np.concatenate(next_state), idxs

    def update_priorities(self, idxs, new_prios):
        for i, p in zip(idxs, new_prios):
            self.priorities[i] = float(abs(p))

    def __len__(self):
        return len(self.buffer)


class SACAgent:
    def __init__(self, state_dim, action_dim):
        self.action_dim = action_dim

        self.actor = Actor(state_dim, action_dim).to(device)
        self.actor_target = Actor(state_dim, action_dim).to(device)
        self.actor_target.load_state_dict(self.actor.state_dict())
        self.actor_optimizer = optim.Adam(self.actor.parameters(), lr=LR_ACTOR)

        self.critic1 = QNetwork(state_dim, action_dim).to(device)
        self.critic2 = QNetwork(state_dim, action_dim).to(device)
        self.critic1_target = QNetwork(state_dim, action_dim).to(device)
        self.critic2_target = QNetwork(state_dim, action_dim).to(device)
        self.critic1_target.load_state_dict(self.critic1.state_dict())
        self.critic2_target.load_state_dict(self.critic2.state_dict())
        self.critic1_optimizer = optim.Adam(self.critic1.parameters(), lr=LR_CRITIC)
        self.critic2_optimizer = optim.Adam(self.critic2.parameters(), lr=LR_CRITIC)

        self.replay_buffer = ReplayMemory(MEMORY_SIZE)
        self.update_count = 0

        self.log_alpha = torch.tensor(float(np.log(0.2)), requires_grad=True, device=device)
        self.alpha_optimizer = optim.Adam([self.log_alpha], lr=5e-4)
        self.target_entropy = np.log(float(self.action_dim))
        self.reward_scale = 1000.0

    @property
    def alpha(self):
        return self.log_alpha.exp()

    def get_action(self, state, add_noise=True):
        state = torch.FloatTensor(state).unsqueeze(0).to(device)
        with torch.no_grad():
            action_probs = self.actor(state)
        dist = torch.distributions.Categorical(action_probs)
        action_index = dist.sample()
        action = [0.0 for _ in range(self.action_dim)]
        action[action_index.item()] = 1.0
        return action

    def update(self):
        if len(self.replay_buffer) < BATCH_SIZE:
            return

        states, actions, rewards, next_states, idxs = self.replay_buffer.sample(BATCH_SIZE)
        states = torch.FloatTensor(states).to(device)
        actions_tmp = np.vstack(actions)
        if actions_tmp.ndim == 2 and actions_tmp.shape[1] == self.action_dim:
            idxs = np.argmax(actions_tmp, axis=1).astype(np.int64)
            action_indices = torch.LongTensor(idxs).to(device)
            actions = F.one_hot(action_indices, num_classes=self.action_dim).float()
        else:
            action_indices = torch.LongTensor(actions_tmp.astype(np.int64)).to(device)
            actions = F.one_hot(action_indices, num_classes=self.action_dim).float()
        rewards = torch.FloatTensor(rewards).unsqueeze(1).to(device)
        rewards_mean = rewards.mean().item()
        self.reward_scale = 0.99 * self.reward_scale + 0.01 * abs(rewards_mean)
        rewards = rewards / (self.reward_scale + 1e-6)
        rewards = torch.clamp(rewards, -10.0, 10.0)
        next_states = torch.FloatTensor(next_states).to(device)

        with torch.no_grad():
            next_action_probs = self.actor_target(next_states)
            q1_next = self.critic1_target(next_states)
            q2_next = self.critic2_target(next_states)
            min_q_next = torch.min(q1_next, q2_next)
            log_next = torch.log(next_action_probs + 1e-8)
            v_next = (next_action_probs * (min_q_next - self.alpha * log_next)).sum(dim=1, keepdim=True)
            target_Q = rewards + GAMMA * v_next

        q1 = (self.critic1(states) * actions).sum(dim=1, keepdim=True)
        q2 = (self.critic2(states) * actions).sum(dim=1, keepdim=True)

        td_err1 = (q1 - target_Q).detach().cpu().numpy().flatten()
        td_err2 = (q2 - target_Q).detach().cpu().numpy().flatten()
        critic1_loss = nn.SmoothL1Loss()(q1, target_Q)
        critic2_loss = nn.SmoothL1Loss()(q2, target_Q)

        self.critic1_optimizer.zero_grad()
        critic1_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.critic1.parameters(), GRAD_CLIP)
        self.critic1_optimizer.step()

        self.critic2_optimizer.zero_grad()
        critic2_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.critic2.parameters(), GRAD_CLIP)
        self.critic2_optimizer.step()

        action_probs = self.actor(states)
        log_probs = torch.log(action_probs + 1e-8)
        with torch.no_grad():
            q1_vals = self.critic1(states)
            q2_vals = self.critic2(states)
            min_q_vals = torch.min(q1_vals, q2_vals)
        actor_loss = (action_probs * (self.alpha * log_probs - min_q_vals)).sum(dim=1).mean()
        self.actor_optimizer.zero_grad()
        actor_loss.backward()
        torch.nn.utils.clip_grad_norm_(self.actor.parameters(), GRAD_CLIP)
        self.actor_optimizer.step()

        target_entropy = self.target_entropy
        with torch.no_grad():
            ap = self.actor(states)
            lp = torch.log(ap + 1e-8)
            entropy_term = (ap * lp).sum(dim=1)
        alpha_loss = -(self.alpha * (entropy_term + target_entropy)).mean()
        self.alpha_optimizer.zero_grad()
        alpha_loss.backward()
        self.alpha_optimizer.step()
        with torch.no_grad():
            self.log_alpha.clamp_(min=-5.0, max=1.0)
            self.target_entropy = max(np.log(float(self.action_dim)) * 0.3, self.target_entropy * ENTROPY_DECAY)

        with torch.no_grad():
            td_mix = 0.5 * (np.abs(td_err1) + np.abs(td_err2))
        self.replay_buffer.update_priorities(idxs, td_mix)

        for target_param, param in zip(self.actor_target.parameters(), self.actor.parameters()):
            target_param.data.copy_(TAU * param.data + (1 - TAU) * target_param.data)
        for target_param, param in zip(self.critic1_target.parameters(), self.critic1.parameters()):
            target_param.data.copy_(TAU * param.data + (1 - TAU) * target_param.data)
        for target_param, param in zip(self.critic2_target.parameters(), self.critic2.parameters()):
            target_param.data.copy_(TAU * param.data + (1 - TAU) * target_param.data)